import pandas as pd
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import text
from indicapp.models import AsyncSessionFactory
from concurrent.futures import ProcessPoolExecutor
from tqdm import tqdm
import numpy as np
import asyncio


# CPU密集型计算函数（单独进程执行）
def calculate_ma(df: pd.DataFrame, ma_periods: list) -> pd.DataFrame:
    """计算移动平均线（在独立进程中运行）"""
    for period in ma_periods:
        df[f'ma{period}'] = df['close'].rolling(window=int(period)).mean().round(2)

    # 只删除 ma2 全部为 NaN 的行，其他行保留
    ma_columns = [f'ma{period}' for period in ma_periods]
    df = df.dropna(subset=ma_columns, how='all')
    
    # 可选：将剩余的 NaN 替换为 None（如果后续需要）
    df = df.replace({np.nan: None})
    return df


# 单个期货在单个周期上的ma计算
async def process_symbol_period(pool: ProcessPoolExecutor, symbol: str, period: str, ma_periods: str):
    """处理单个task的MA计算"""
    session = AsyncSessionFactory() # 每一个协程，需要有独属于自己的session
    try:
        async with session.begin():
            # 获取K线数据
            table_name = f"{symbol.lower()}_{period}min"
            kline_data = await session.execute(
                text(f"SELECT datetime, close FROM klinedb.{table_name} ORDER BY datetime")
            )
            df = pd.DataFrame(kline_data.fetchall(), columns=['datetime', 'close'])

            loop = asyncio.get_running_loop()
            result_df = await loop.run_in_executor(
                pool, 
                calculate_ma, 
                df.copy(),  # 避免进程间共享内存
                ma_periods  # MA周期
            )

            # 存储结果
            flag = await store_ma_results(session, symbol, period, ma_periods, result_df)
            return symbol, period, flag
    except Exception as e:
        tqdm.write(f"合约{symbol}在{period}min上的均线MA处理失败: {str(e)}") 
        return symbol, period, False


async def store_ma_results(session: AsyncSession, symbol: str, period: str, ma_periods: list, df: pd.DataFrame):
    temp1 = "datetime"
    temp2 = "datetime DATETIME PRIMARY KEY"
    temp3 = ":datetime"
    for ma_period in ma_periods:
        temp1 += f", ma{ma_period}"
        temp2 += f", ma{ma_period} FLOAT"
        temp3 += f", :ma{ma_period}"
    
    """存储计算结果到数据库"""
    try:
        await session.execute(
            text(f"CREATE TABLE IF NOT EXISTS ma_{symbol.lower()}_{period}min ({temp2})")
        )
    except Exception as e:
        tqdm.write(f"为{symbol}创建{period}min均线表失败：{str(e)}")
        return False

    try:
        await session.execute(
            text(f"""
            INSERT INTO ma_{symbol.lower()}_{period}min ({temp1})
            VALUES ({temp3})
            """),
            [dict(row) for _, row in df.iterrows()]
        )
    except Exception as e:
        tqdm.write(f"{symbol}插入{period}min均线数据失败: {str(e)}") 
        return False
    return True


async def process_tasks_concurrently(symbols, periods, ma_periods):
    # 使用多进程计算MA（一个核心开一个进程，我的CPU是2核，所以我起2个进程）
    with ProcessPoolExecutor(max_workers=2) as pool:
        success_count = 0  # 成功计数器
        # 创建所有任务
        tasks = [
            # 这里直接选择将pool作为参数传给任务，这样tasks里面的所有任务都共用一个线程池
            process_symbol_period(pool, symbol, period, ma_periods)
            for symbol in symbols
            for period in periods
        ]
        total_tasks = len(tasks)  # 保存总任务数
        # 并发执行并显示进度
        with tqdm(total=total_tasks, desc="均线MA计算/入库进度") as pbar:
            for task in asyncio.as_completed(tasks):
                symbol, period, result = await task
                if result is True:
                    success_count += 1
                    pbar.update(1)
                    pbar.set_postfix_str(f"完成: {symbol}_{period}min")
                else:
                    pbar.set_postfix_str(f"失败: {symbol}_{period}min")
                pbar.refresh()

    if success_count != total_tasks:
        failed_count = total_tasks - success_count
        error_msg = (f"部分MA计算未完成，失败: {failed_count}/{total_tasks}")
        raise RuntimeError(error_msg)